import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import losses, optimizers, metrics, layers, activations, utils
import os

VER = '2.1'
N_CLS = 10
BATCH_SIZE = 128
N_EPOCH = 8
ALPHA = 0.001
SPEC = VER + '_' + str(ALPHA) + '_' + str(N_EPOCH) + '_' + str(BATCH_SIZE)
FILE_NAME = os.path.basename(__file__)
SAVE_DIR = os.path.join('_save', FILE_NAME, SPEC)
SAVE_PREFIX = os.path.join(SAVE_DIR, 'weights')

input = keras.Input((28, 28, 1))

x = layers.Conv2D(6, (5, 5))(input)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
x = layers.MaxPool2D(strides=(2, 2), padding='same')(x)

x = layers.Conv2D(16, (3, 3))(x)
x = layers.BatchNormalization()(x)
x = layers.ReLU()(x)
x = layers.MaxPool2D(strides=(2, 2), padding='same')(x)

x = layers.Flatten()(x)
x = layers.Dense(120, activation=activations.relu)(x)
x = layers.Dense(84, activation=activations.relu)(x)

output = layers.Dense(N_CLS, activation=activations.softmax)(x)
model = keras.Model(input, output)

model.compile(
    loss=losses.categorical_crossentropy,
    optimizer=optimizers.Adam(lr=ALPHA),
    metrics=metrics.categorical_accuracy
)

if '__main__' == __name__:

    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
    print(x_train.shape, y_train.shape)  # 60000x28x28
    print(x_test.shape, y_test.shape)

    x_train = np.expand_dims(x_train, 3)
    x_test = np.expand_dims(x_test, 3)
    x_train = x_train.astype(np.float32) / 255.
    x_test = x_test.astype(np.float32) / 255.
    y_train = utils.to_categorical(y_train, N_CLS)
    y_test = utils.to_categorical(y_test, N_CLS)
    print(x_train.shape, y_train.shape)  # 60000x28x28x1
    print(x_test.shape, y_test.shape)

    if os.path.exists(SAVE_DIR):
        print('Loading...')
        model.load_weights(SAVE_PREFIX)
        print('Loaded')
    else:
        model.fit(x_train, y_train,
                  batch_size=BATCH_SIZE, epochs=N_EPOCH)
        os.makedirs(SAVE_DIR, exist_ok=True)
        model.save_weights(SAVE_PREFIX)
        print('Saved')

    r = model.evaluate(x_test, y_test,
                       batch_size=BATCH_SIZE)
    print(r)
